#include <iostream>
#include <vector>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>

typedef std::vector<std::vector<uchar>> MatrixU;
typedef std::vector<std::vector<double>> MatrixD;

enum BORDER_TYPE {
    BORDER_CONSTANT = 0,
    BORDER_REPLICATE,
    BORDER_REFLECT,
    BORDER_WRAP,
    BORDER_REFLECT101
};

MatrixU copy_make_border(const MatrixU& mat, int top, int bottom, int left, int right, int borderType, uchar val=0) {
    int rows = mat.size();
    int cols = 0;
    if (rows > 0) {
        cols = mat[0].size();
    }
    if (cols == 0) {
        assert(false);
    }
    
    MatrixU result;
    int result_rows = rows + top + bottom;
    int result_cols = cols + left + right;
    int i, j;
    for (i = 0; i < result_rows; i++) {
        std::vector<uchar> v;
        v.resize(result_cols);
        result.push_back(v);
    }
    
    // Copy mat to center of result
    for (i = 0; i < rows; i++) {
        for (j = 0; j < cols; j++) {
            result[top+i][left+j] = mat[i][j];
        }
    }
    
    switch (borderType) {
        case BORDER_CONSTANT:
            for (i = 0; i < result_rows; i++) {
                for (j = 0; j < result_cols; j++) {
                    if (i < top || i >= top+rows) {  // Top, Bottom
                        result[i][j] = val;
                    } else {
                        if (j < left || j >= left+cols) {  // Left, Right
                            result[i][j] = val;
                        }
                    }
                }
            }
            break;
        case BORDER_REPLICATE:
            // Top
            for (i = 0; i < top; i++) {
                for (j = left; j < left+cols; j++) {
                    result[i][j] = mat[0][j-left];
                }
            }
            // Bottom
            for (i = top+rows; i < result_rows; i++) {
                for (j = left; j < left+cols; j++) {
                    result[i][j] = mat[rows-1][j-left];
                }
            }
            // Left
            for (i = 0; i < result_rows; i++) {
                for (j = 0; j < left; j++) {
                    result[i][j] = result[i][left];
                }
            }
            // Right
            for (i = 0; i < result_rows; i++) {
                for (j = left+cols; j < result_cols; j++) {
                    result[i][j] = result[i][left+cols-1];
                }
            }
            break;
        case BORDER_REFLECT:
            // Top
            for (i = 0; i < top; i++) {
                for (j = left; j < left+cols; j++) {
                    result[i][j] = mat[top-1-i][j-left];
                }
            }
            // Bottom
            for (i = top+rows; i < result_rows; i++) {
                for (j = left; j < left+cols; j++) {
                    result[i][j] = mat[rows-1-(i-top-rows)][j-left];
                }
            }
            // Left
            for (i = 0; i < result_rows; i++) {
                for (j = 0; j < left; j++) {
                    result[i][j] = result[i][left+left-1-j];
                }
            }
            // Right
            for (i = 0; i < result_rows; i++) {
                for (j = left+cols; j < result_cols; j++) {
                    result[i][j] = result[i][left+cols-1-(j-left-cols)];
                }
            }
            break;
        case BORDER_WRAP:
            // Top
            for (i = 0; i < top; i++) {
                for (j = left; j < left+cols; j++) {
                    result[i][j] = mat[rows-top+i][j-left];
                }
            }
            // Bottom
            for (i = top+rows; i < result_rows; i++) {
                for (j = left; j < left+cols; j++) {
                    result[i][j] = mat[i-top-rows][j-left];
                }
            }
            // Left
            for (i = 0; i < result_rows; i++) {
                for (j = 0; j < left; j++) {
                    result[i][j] = result[i][cols+j];
                }
            }
            // Right
            for (i = 0; i < result_rows; i++) {
                for (j = left+cols; j < result_cols; j++) {
                    result[i][j] = result[i][j-cols];
                }
            }
            break;
        case BORDER_REFLECT101:
            // Top
            for (i = 0; i < top; i++) {
                for (j = left; j < left+cols; j++) {
                    result[i][j] = mat[top-i][j-left];
                }
            }
            // Bottom
            for (i = top+rows; i < result_rows; i++) {
                for (j = left; j < left+cols; j++) {
                    result[i][j] = mat[rows-2-(i-top-rows)][j-left];
                }
            }
            // Left
            for (i = 0; i < result_rows; i++) {
                for (j = 0; j < left; j++) {
                    result[i][j] = result[i][left+left-j];
                }
            }
            // Right
            for (i = 0; i < result_rows; i++) {
                for (j = left+cols; j < result_cols; j++) {
                    result[i][j] = result[i][left+cols-2-(j-left-cols)];
                }
            }
            break;
        default:
            assert(false);
    }
    return result;
}

MatrixU convolution(const MatrixU& matPadded, const MatrixD& kernel) {
    int padded_rows = matPadded.size();
    int padded_cols = matPadded[0].size();
    int krows = kernel.size();
    int kcols = kernel[0].size();
    int top = krows / 2;
    int bottom = krows / 2;
    int left = kcols / 2;
    int right = kcols / 2;
    int rows = padded_rows - top - bottom;
    int cols = padded_cols - left - right;
    
    MatrixU convolve;
    for (int i = 0; i < rows; i++) {
        std::vector<uchar> v;
        v.resize(cols);
        convolve.push_back(v);
    }
    
    for (int i = 0; i < rows; i++) {
        for (int j = 0; j < cols; j++) {
            double s = 0;
            for (int ki = 0; ki < krows; ki++) {
                for (int kj = 0; kj < kcols; kj++) {
                    s += kernel[ki][kj] * matPadded[i+ki][j+kj];
                }
            }
            if (s < 0) {
                s = 0;
            }
            if (s > 255) {
                s = 255;
            }
            convolve[i][j] = (uchar)std::round(s);
        }
    }
    return convolve;
}

#define GAUSSIAN_KERNEL_SIZE 5
MatrixD get_gaussian_kernel() {
    MatrixD gk;
    std::vector<double> r1 = {1, 4, 6, 4, 1},
    r2 = {4, 16, 24, 16, 4},
    r3 = {6, 24, 36, 24, 6},
    r4 = {4, 16, 24, 16, 4},
    r5 = {1, 4, 6, 4, 1};
    for (int i = 0; i < GAUSSIAN_KERNEL_SIZE; i++) {
        r1[i] /= 256.;
        r2[i] /= 256.;
        r3[i] /= 256.;
        r4[i] /= 256.;
        r5[i] /= 256.;
    }
    gk.push_back(r1);
    gk.push_back(r2);
    gk.push_back(r3);
    gk.push_back(r4);
    gk.push_back(r5);
    return gk;
}

void copy_data_from_cv_mat_to_matrix(const cv::Mat& src, MatrixU& dst) {
    int rows = src.rows;
    int cols = src.cols;
    dst.resize(rows);
    for (int i = 0; i < rows; i++) {
        const uchar* p = src.ptr<uchar>(i);
        std::vector<uchar> v;
        for (int j = 0; j < cols; j++) {
            v.push_back(p[j]);
        }
        dst[i] = v;
    }
}

void copy_data_from_matrix_to_cv_mat(const MatrixU& src, cv::Mat& dst) {
    int rows = src.size();
    int cols = src[0].size();
    dst = cv::Mat(rows, cols, CV_8UC1);
    for (int i = 0; i < rows; i++) {
        uchar* p = dst.ptr<uchar>(i);
        for (int j = 0; j < cols; j++) {
            p[j] = src[i][j];
        }
    }
}
void copy_data_from_matrix_to_cv_mat(const MatrixD& src, cv::Mat& dst) {
    int rows = src.size();
    int cols = src[0].size();
    dst = cv::Mat(rows, cols, CV_64FC1);
    for (int i = 0; i < rows; i++) {
        double* p = dst.ptr<double>(i);
        for (int j = 0; j < cols; j++) {
            p[j] = src[i][j];
        }
    }
}

MatrixU gaussian_pyramid_downsampling(const MatrixU& mat, int borderType) {
    MatrixD kernel;
    kernel = get_gaussian_kernel();
    MatrixU matPadded;
    int top = GAUSSIAN_KERNEL_SIZE / 2;
    int bottom = GAUSSIAN_KERNEL_SIZE / 2;
    int left = GAUSSIAN_KERNEL_SIZE / 2;
    int right = GAUSSIAN_KERNEL_SIZE / 2;
    matPadded = copy_make_border(mat, top, bottom, left, right, borderType, 0);
    
    MatrixU convolve;
    convolve = convolution(matPadded, kernel);
    
    int rows = mat.size();
    int cols = mat[0].size();
    int result_rows = (rows+1)/2;
    int result_cols = (cols+1)/2;
    MatrixU result;
    for (int i = 0; i < result_rows; i++) {
        std::vector<uchar> v;
        v.resize(result_cols);
        result.push_back(v);
    }
    for (int i = 0; i < result_rows; i++) {
        for (int j = 0; j < result_cols; j++) {
            result[i][j] = convolve[i*2][j*2];
        }
    }
    return result;
}

MatrixU laplacian_pyramid_upsampling(const MatrixU& mat, int borderType) {
    int rows = mat.size();
    int cols = mat[0].size();
    MatrixU res;
    for (int i = 0; i < rows*2; i++) {
        std::vector<uchar> v;
        v.resize(cols*2);
        for (int j = 0; j < cols*2; j++) {
            v[j] = 0;
        }
        res.push_back(v);
    }
    
    for (int i = 0; i < rows; i++) {
        for (int j = 0; j < cols; j++) {
            res[i*2][j*2] = mat[i][j];
        }
    }
    
    MatrixD kernel;
    kernel = get_gaussian_kernel();
    for (int i = 0; i < kernel.size(); i++) {
        for (int j = 0; j < kernel[i].size(); j++) {
            kernel[i][j] *= 4;
        }
    }
    
    MatrixU resPadded;
    int top = GAUSSIAN_KERNEL_SIZE / 2;
    int bottom = GAUSSIAN_KERNEL_SIZE / 2;
    int left = GAUSSIAN_KERNEL_SIZE / 2;
    int right = GAUSSIAN_KERNEL_SIZE / 2;
    resPadded = copy_make_border(res, top, bottom, left, right, borderType, 0);
    
    res = convolution(resPadded, kernel);
    return res;
}

int main(int argc, char **argv) {
    cv::Mat lena = cv::imread("/home/xlll/Downloads/opencv/samples/data/lena.jpg", cv::IMREAD_GRAYSCALE);
    if (lena.empty()) {
        std::cout << "failed to read image" << std::endl;
        return EXIT_FAILURE;
    }
    
    int rows = lena.rows;
    int cols = lena.cols;
    MatrixU mat;
    copy_data_from_cv_mat_to_matrix(lena, mat);
    int borderType = BORDER_REFLECT101;
    
    cv::Mat downsample1, downsample2;
    cv::pyrDown(lena, downsample1, cv::Size(), borderType);
    MatrixU downsampleResult;
    downsampleResult = gaussian_pyramid_downsampling(mat, borderType);
    copy_data_from_matrix_to_cv_mat(downsampleResult, downsample2);
    cv::Mat diff1 = downsample1 != downsample2;
    std::vector<cv::Point> nonzero1;
    cv::findNonZero(diff1, nonzero1);
    std::cout << "nonzero1 size = " << nonzero1.size() << std::endl;
    cv::imshow("lena", lena);
    cv::imshow("downsample1", downsample1);
    cv::imshow("downsample2", downsample2);
    
    cv::Mat upsample1, upsample2;
    cv::pyrUp(lena, upsample1, cv::Size(), borderType);
    MatrixU upsampleResult;
    upsampleResult = laplacian_pyramid_upsampling(mat, borderType);
    copy_data_from_matrix_to_cv_mat(upsampleResult, upsample2);
    cv::Mat diff2 = upsample1 != upsample2;
    std::vector<cv::Point> nonzero2;
    cv::findNonZero(diff2, nonzero2);
    std::cout << "nonzero2 size = " << nonzero2.size() << std::endl;
    cv::imshow("upsample1", upsample1);
    cv::imshow("upsample2", upsample2);
    
    cv::waitKey(0);
    return 0;
}
